#include "Encoder.h"
#include <fstream>
#include <stdlib.h>
#include "Decoder.h"
#include <algorithm>
#include "nvToolsExt.h"

using namespace std;

Encoder::Encoder(QuantizationMatrix* qm, bool rle, int width, int height,  int matrixColumnSize)
{
    /*Size*/
    _matrixColumnLength = matrixColumnSize;
    _matrixSize = matrixColumnSize*matrixColumnSize;

    /*Matrix*/
    _DCTMatrix = new Matrix(_matrixColumnLength);
    _tDCTMatrix = new Matrix(_matrixColumnLength);
    initDCTMatrices();
    _quantizationMatrix = qm;

    /*Vars*/
    _RLE = rle;
    _width = width;
    _height = height;
    _columnSize = matrixColumnSize;

    /*BlockList*/
    _totalBlocks = (_width/matrixColumnSize)*(_height/matrixColumnSize);
}

Encoder::Encoder(EncoderSettings &config, QuantizationMatrix* qm, int matrixColumnSize)
{
    /*Size*/
    _matrixColumnLength = matrixColumnSize;
    _matrixSize = matrixColumnSize*matrixColumnSize;

    /*Matrix*/
    _DCTMatrix = new Matrix(_matrixColumnLength);
    _tDCTMatrix = new Matrix(_matrixColumnLength);
    initDCTMatrices();
    _quantizationMatrix = qm;

    /*Vars*/
    _RLE = config.getRLE();
    _width = config.getWidth();
    _height = config.getHeight();
    _columnSize = matrixColumnSize;
    _meRange = config.getMerange();

    /*BlockList*/
    _totalBlocks = (_width/matrixColumnSize)*(_height/matrixColumnSize);

    /*File names*/
    _logFile = config.logfile();
}

util::BitStreamWriter* Encoder::encodePFrame(uint8_t *buff, uint8_t* previousFrame)
{
	nvtxEventAttributes_t initAttrib = {0};
	initAttrib.version = NVTX_VERSION;
	initAttrib.size = NVTX_EVENT_ATTRIB_STRUCT_SIZE;
	initAttrib.color = 0xFF880000;
	initAttrib.colorType = NVTX_COLOR_ARGB;
	initAttrib.message.ascii = "PFrame";
	initAttrib.messageType = NVTX_MESSAGE_TYPE_ASCII;
	initAttrib.category = 1;
	nvtxRangePushEx(&initAttrib);
    _blockList.initialize(buff, _width, _height, _totalBlocks, _columnSize);

    int size = _totalBlocks*_matrixSize;

    _blockList.it = _blockList.begin();

    uint8_t* buffer = new uint8_t[size];

    util::BitStreamWriter* bsw = new util::BitStreamWriter(buffer, size);

    util::BitStreamWriter** bswArray = new util::BitStreamWriter*[_totalBlocks];

#if P_ENC_PARALLEL == 1
#pragma omp parallel for
#endif
    for(int i = 0; i < _blockList.size(); ++i)
    {
        Matrix m = _blockList.get(i);
        int8_t vectorX;
        int8_t vectorY;
        Matrix* bestComparingBlock;
        Matrix* diff;
        util::BitStreamWriter* blockWriter = new util::BitStreamWriter(_matrixSize);
        bswArray[i] = blockWriter;

        uint32_t type = calculateMotionVector(&vectorX, &vectorY, &bestComparingBlock, &diff, &m, ((i*_columnSize)%_width), (((i*_columnSize)/_width)*_columnSize), previousFrame);

        blockWriter->put(2, type);
        if(type == 0)
        {
            if(vectorX < 0)
                blockWriter->put_bit(1);
            else
                blockWriter->put_bit(0);
            blockWriter->put(7, (uint32_t)vectorX);
            if(vectorY < 0)
                blockWriter->put_bit(1);
            else
                blockWriter->put_bit(0);
            blockWriter->put(7, (uint32_t)vectorY);
        }
        else if(type == 1)
        {
            blockWriter->put(8, (uint32_t)vectorX);
            blockWriter->put(8, (uint32_t)vectorY);
            Encoder matrixEncoder(_quantizationMatrix, _RLE, 16, 16, 4);
            uint8_t* macroBlockData = diff->getMatrixData();
            util::BitStreamWriter* encodedMacroBlock = matrixEncoder.encode(macroBlockData);
            blockWriter->put_bit(0);
            uint32_t bytesWritten = (uint32_t)(encodedMacroBlock->get_position()/8);
            if(encodedMacroBlock->get_position()%8 != 0)
                ++bytesWritten;
            blockWriter->put(31, bytesWritten);
            util::BitStreamReader encodedMacroBlockReader(encodedMacroBlock->get_buffer(), encodedMacroBlock->get_size());
            for(int k = 0; k < encodedMacroBlock->get_position(); ++k)
                blockWriter->put_bit(encodedMacroBlockReader.get_bit());
        }
        else
        {
            Encoder matrixEncoder(_quantizationMatrix, _RLE, 16, 16, 4);
            uint8_t* macroBlockData = diff->getMatrixData();
            util::BitStreamWriter* encodedMacroBlock  = matrixEncoder.encode(macroBlockData);
            blockWriter->put_bit(0);
            uint32_t bytesWritten = (uint32_t)(encodedMacroBlock->get_position()/8);
            if(encodedMacroBlock->get_position()%8 != 0)
                ++bytesWritten;
            blockWriter->put(31, bytesWritten);
            util::BitStreamReader encodedMacroBlockReader(encodedMacroBlock->get_buffer(), encodedMacroBlock->get_size());
            for(int k = 0; k < encodedMacroBlock->get_position(); ++k)
                blockWriter->put_bit(encodedMacroBlockReader.get_bit());
        }
    }

    for(int i = 0; i < _totalBlocks; ++i)
    {
        util::BitStreamWriter* temp = bswArray[i];

util:BitStreamReader bsr(temp->get_buffer(), temp->get_size());
        for(int j = 0; j < temp->get_position(); ++j)
            bsw->put_bit(bsr.get_bit());
    }

	nvtxRangePop();

    return bsw;
}

struct MADInfoContainer {
    Matrix diff;
    Matrix comparingBlock;
    double MAD;
    int8_t x;
    int8_t y;
};

uint32_t Encoder::calculateMotionVector(int8_t *x, int8_t *y, Matrix** bestComparingBlock, Matrix** diff, Matrix *targetBlock, int tbX, int tbY, uint8_t* previousFrame)
{
    int type = 2;

    *diff = NULL;
    *bestComparingBlock = NULL;
    double lowestMAD = INT_MAX;

    int centerX = tbX;
    int centerY = tbY;
    int lowX = tbX-(_meRange/2);
    if(lowX < 0)
        lowX = 0;
    int lowY = tbY-(_meRange/2);
    if(lowY < 0)
        lowY = 0;
    int highX = tbX+(_meRange/2);
    int highY = tbY+(_meRange/2);

    int limX = min(highX, _width);
    int limY = min(highY, _height);

    int columns = limX - lowX;
    int rows = limY - lowY;
    int size = rows*columns;

    MADInfoContainer* MADArray = new MADInfoContainer[size];

#if P_ENC_ME == 1
#pragma omp parallel for
#endif
    for(int k = 0; k < size; ++k)
    {
        int i = k / rows + lowX;
        int j = k % rows + lowY;
        Matrix* comparingBlock = NULL;

        comparingBlock = BlockList::getBlockAtCoord(i, j, previousFrame, _width, _height, _columnSize);

        MADInfoContainer& MIC = MADArray[k];

        MIC.diff  = (*targetBlock) - (*comparingBlock);

        double currentMAD = calculateMAD(MIC.diff);

        MIC.MAD = currentMAD;

        MIC.comparingBlock = *comparingBlock;
        MIC.x = i-centerX;
        MIC.y = j-centerY;

        delete comparingBlock;
    }

    for(int i = 0; i < size; ++i)
    {
        if(MADArray[i].MAD < lowestMAD)
        {
            lowestMAD = MADArray[i].MAD;
            *x = MADArray[i].x;
            *y = MADArray[i].y;
            if(*diff != NULL)
                delete *diff;
            *diff = new Matrix(_columnSize);
            **diff = MADArray[i].diff;
            if(*bestComparingBlock != NULL)
                delete *bestComparingBlock;
            *bestComparingBlock = new Matrix(_columnSize);
            **bestComparingBlock = MADArray[i].comparingBlock;
            type = 1;
        }
    }

    delete[] MADArray;

    if(lowestMAD == 0)
        type = 0;

    return type;

}

Matrix* Encoder::calculateDifference(Matrix *target, Matrix *reference)
{
    Matrix *diff = new Matrix(_columnSize);

    *diff = (*target) - (*reference);

    return diff;
}

double Encoder::calculateMAD(Matrix& matrix)
{
    int madTot = 0;

    for(int i = 0; i < _columnSize*_columnSize; ++i)
    {
        int val = matrix[i];
        if(val < 0)
            val = -val;

        madTot += val;

    }

    return (double)madTot/16.0;
}

util::BitStreamWriter* Encoder::encode(uint8_t* buff)
{
	nvtxRangePushA("Encode IFrame");
    _blockList.initialize(buff, _width, _height, _totalBlocks, _columnSize);

    int size = _totalBlocks*_matrixSize;

#if I_ENC_PARALLEL == 1
#pragma omp parallel for
#endif

    for(int i = 0; i < _blockList.size(); ++i)
    {
        Matrix& m = _blockList.get(i);
        performNormalization(m);
        performFDCT(m);
        performQuantization(m);
    }


    int8_t* zigzag;
    zigzag = performZigZag();

    uint8_t* buffer = new uint8_t[size];
    util::BitStreamWriter* bsw = new util::BitStreamWriter(buffer, size);

    if(_RLE)
        performRLEncoding(zigzag, bsw);
    else
        performRawEncoding(zigzag, bsw);

    delete[] zigzag;

    if(bsw->get_position()%8 != 0)
        bsw->flush();

	nvtxRangePop();

    return bsw;
}

void Encoder::initDCTMatrices()
{
    double dct[16];
    double tdct[16];

    dct[ 0] =  0.5;
    dct[ 1] =  0.5;
    dct[ 2] =  0.5;
    dct[ 3] =  0.5;
    dct[ 4] =  0.653;
    dct[ 5] =  0.271;
    dct[ 6] = -0.271;
    dct[ 7] = -0.653;
    dct[ 8] =  0.5;
    dct[ 9] = -0.5;
    dct[10] = -0.5;
    dct[11] =  0.5;
    dct[12] =  0.271;
    dct[13] = -0.653;
    dct[14] =  0.653;
    dct[15] = -0.271;

    _DCTMatrix->setMatrixRowMajor(dct);

    tdct[ 0] =  0.5;
    tdct[ 1] =  0.653;
    tdct[ 2] =  0.5;
    tdct[ 3] =  0.271;
    tdct[ 4] =  0.5;
    tdct[ 5] =  0.271;
    tdct[ 6] = -0.5;
    tdct[ 7] = -0.653;
    tdct[ 8] =  0.5;
    tdct[ 9] = -0.271;
    tdct[10] = -0.5;
    tdct[11] =  0.653;
    tdct[12] =  0.5;
    tdct[13] = -0.653;
    tdct[14] =  0.5;
    tdct[15] = -0.271;

    _tDCTMatrix->setMatrixRowMajor(tdct);
}

void Encoder::performNormalization(Matrix& m)
{
    m.normalize();
}

void Encoder::performFDCT(Matrix& m)
{
    m = (*_DCTMatrix) * m * (*_tDCTMatrix);
    m.round();

}

void Encoder::performQuantization(Matrix& m)
{
    m = m / *_quantizationMatrix;
    m.flattenHighValues();
}

int8_t *Encoder::performZigZag()
{
    int8_t *imageList = new int8_t[_totalBlocks*16];
    int position=0;

    for(_blockList.it = _blockList.begin(); _blockList.it != _blockList.end(); ++_blockList.it){
        int8_t *temp = (*_blockList.it).zigzag();

        for(int i=0; i < 16; ++i)
            imageList[position++] = temp[i];

        delete[] temp;
    }

    return imageList;
}

unsigned long Encoder::performRLEncoding(int8_t *imageList, util::BitStreamWriter *bsw)
{
    int seriesOf0=0;
    unsigned long bitsWritten = 0;
    for(int i=0; i < _totalBlocks; ++i) {
        for(int j=0; j < 16; ++j) {
            int8_t currentNumber = imageList[(i*16)+j];
            if (currentNumber != 0) {

                if(seriesOf0)       // if any 0 on queue
                    bitsWritten += fillLast0Series(bsw, seriesOf0);

                bitsWritten += fillNumber(bsw, currentNumber);

            } else {
                ++seriesOf0;
            }
        }
        if(seriesOf0)       // if any 0 on queue
            bitsWritten += fillLast0Series(bsw, seriesOf0);
    }
    return bitsWritten;
}

unsigned long Encoder::performRawEncoding(int8_t *imageList, util::BitStreamWriter *bsw)
{
    unsigned long bitsWritten = 0;
    for(int i=0; i < _totalBlocks; ++i)
        for(int j=0; j < 16; ++j)
            bitsWritten += fillNumber(bsw, imageList[(i*16)+j]);
    return bitsWritten;
}

unsigned long Encoder::fillLast0Series(util::BitStreamWriter *bsw, int &seriesOf0)
{
    unsigned long bitsWritten = 0;
    bsw->put_bit(1);
    ++bitsWritten;

    --seriesOf0;

    bitsWritten+=4;
    bsw->put(4, seriesOf0);

    seriesOf0 = 0;

    return bitsWritten;
}

unsigned long Encoder::fillNumber(util::BitStreamWriter *bsw, int8_t &currentNumber)
{
    uint32_t numberOfBits = 0;
    unsigned long bitsWritten = 0;

    if(currentNumber < 2 && currentNumber > -3)
    {
        numberOfBits = 1;
    }
    else if(currentNumber < 4 && currentNumber > -5)
    {
        numberOfBits = 2;
    }
    else if(currentNumber < 8 && currentNumber > -9)
    {
        numberOfBits = 3;
    }
    else if(currentNumber < 16 && currentNumber > -17)
    {
        numberOfBits = 4;
    }
    else if(currentNumber < 32 && currentNumber > -33)
    {
        numberOfBits = 5;
    }
    else if(currentNumber < 64 && currentNumber > -65)
    {
        numberOfBits = 6;
    }
    else if(currentNumber < 128 && currentNumber > -129)
    {
        numberOfBits = 7;
    }

    if(numberOfBits > 0){
        bsw->put_bit(0);
        ++bitsWritten;
        bsw->put(3, numberOfBits-1);
        bitsWritten+=3;
        if(currentNumber >= 0)
            bsw->put_bit(0);
        else
        {
            bsw->put_bit(1);
            currentNumber = (-1*currentNumber)-1;
        }
        bsw->put(numberOfBits, currentNumber);
        bitsWritten+=numberOfBits;
    }

    return bitsWritten;
}
